from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain

import gradio as gr


llm = ChatOpenAI(temperature=0)


def initialize_sales_data():
    embeddings = OpenAIEmbeddings()
    try:
        FAISS.load_local(folder_path="./faiss_sales",
                         embeddings=embeddings,
                         index_name="sales_index",
                         allow_dangerous_deserialization=True)
    except RuntimeError:

        loader = TextLoader("real_estate_sales_data.txt")
        documents = loader.load()
        text_splitter = CharacterTextSplitter(
            separator=r'\d+\.',
            chunk_size=100,
            chunk_overlap=0,
            length_function=len,
            is_separator_regex=True,
        )
        docs = text_splitter.split_documents(documents=documents)
        embeddings = OpenAIEmbeddings()
        db = FAISS.from_documents(documents=docs, embedding=embeddings)
        db.save_local(folder_path="./faiss_sales", index_name="sales_index")
    # retriever = db.as_retriever(search_type="similarity_score_threshold",
    #                             search_kwargs={'score_threshold': 0.8, 'k': 1})
    # # results = retriever.get_relevant_documents("这个小区交通便利吗？")
    # # print(results)
    # results = retriever.invoke("这个小区交通便利吗？")
    # print(results)
    # for doc in results:
    #     print(doc.page_content + "\n")
    #
    # result = results[0].page_content.split("[销售回答] ")
    # print(result[-1])


def sales_chat(message, history):
    print(history)
    embeddings = OpenAIEmbeddings()
    db = FAISS.load_local(folder_path="./faiss_sales",
                          embeddings=embeddings,
                          index_name="sales_index",
                          allow_dangerous_deserialization=True)
    retriever = db.as_retriever(search_type="similarity_score_threshold",
                                search_kwargs={'score_threshold': 0.8, 'k': 1})
    results = retriever.invoke(message)
    # print(results)
    # print(type(results))
    if not results:
        return "超出我能力范围，请换个问题提问。"
    chat_prompt = ChatPromptTemplate.from_template("""
        根据下面内容回答问题，回答的问题要根据内容详细描述，不需要超出内容的回复，并且直接返回答案不需要按“回答：”格式：
        
        {context}
        
        问题：{question}
    """)
    # document_chain = create_stuff_documents_chain(llm, chat_prompt)
    # document_chain.input_schema()
    # retriever_chain = create_retrieval_chain(retriever=retriever, combine_docs_chain=document_chain)
    # retriever_chain.input_schema()
    chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | chat_prompt
            | llm
            | StrOutputParser()
    )
    results = chain.invoke(message)
    if not results:
        return "超出我能力范围，请换个问题提问。"
    # results = retriever.invoke(message)
    # print(results)
    # for doc in results:
    #     print(doc.page_content + "\n")
    #
    # result = results[0].page_content.split("[销售回答] ")
    return results


def launch_gradio():
    demo = gr.ChatInterface(fn=sales_chat,
                            title="机器人",
                            chatbot=gr.Chatbot(height=600))
    demo.launch(share=True)


if __name__ == '__main__':
    initialize_sales_data()
    # sales_chat("哪个小区交通便利？")
    launch_gradio()
